import sys
import numpy as np
import torch
import torch.nn as nn
import ctypes

sys.path.append("../utils")

from utils import OpRunner, compare, verify_result
import time


def main():
    func = OpRunner("./build/libop.so", 3, pint_time=True)
    func.op.run.argtypes[2] = ctypes.c_char_p
    minvalue = 1
    maxvalue = 10
    shape = [1024]
    for dtype in [np.float16, np.float32]:
        print("测试类型", dtype)
        for reduction in ["mean", "sum", "none"]:
            input_predict = np.random.uniform(minvalue, maxvalue, shape).astype(dtype)
            input_label = np.random.uniform(minvalue, maxvalue, shape).astype(dtype)
            output = np.zeros(shape, dtype=dtype)
            # 初始化损失函数
            mse_loss = nn.MSELoss(reduction=reduction)
            # 假设我们有一些预测值和目标值
            predictions = torch.tensor(input_predict)
            targets = torch.tensor(input_label)
            # 计算损失
            loss: torch.Tensor = mse_loss(predictions, targets)
            golden = loss.numpy()

            func([input_predict, input_label, reduction.encode(), output], [3])

            print("output", output[0])
            print("golden", golden)
            if reduction == "none":
                verify_result(output, golden)
            else:
                verify_result(np.array(output[0], dtype=dtype), np.array(golden, dtype=dtype))

            print("----------------------------")


if __name__ == "__main__":
    main()
